import galaxy_vis
import random
import math
import Image, ImageDraw
import ImageFilter

def find(f, seq):
  """Return first item in sequence where f(item) == True."""
  for item in seq:
    if f(item): 
      return item

def calc_gaussian(sigma):
  """Calculate 5x5 gauss kernel"""
  w = -2
  h = -2
  k = []
  while h<3:
    w = -2
    s = []
    while w<3:
      k.append((1/(2*math.pi*sigma**2))*math.exp(-((w)**2+(h)**2)/(2*sigma**2)))
      w += 1
    h += 1
  return k

def branch_gen(arg, xoffset, yoffset, a, angle, colorstart, fillstep, widthstart, widthstep):
  im = arg[0]
  draw = arg[1]

  # r^2 = a^2*phi
  phi = 0
  step = 0.1
  c = 1
  prevptx = xoffset
  prevpty = yoffset
  fill = colorstart
  width = widthstart
  while True:
    r = a*math.sqrt(math.fabs(phi))
    ptx = -r*math.cos(phi+angle)+xoffset
    pty = r*math.sin(phi+angle)+yoffset
    draw.line((int(prevptx), int(prevpty), int(ptx), int(pty)), fill=int(fill), width=int(width))
    width += widthstep
    fill += fillstep if fill < 255 else 0
    prevptx = ptx
    prevpty = pty
    if ptx < 0 or pty < 0 or ptx > im.size[0] or pty > im.size[1]:
      break
    phi+=step

  phi = 0
  fill = colorstart
  width = widthstart
  prevptx = xoffset
  prevpty = yoffset
  while True:
    r = -a*math.sqrt(math.fabs(phi))
    ptx = -r*math.cos(phi+angle)+xoffset
    pty = r*math.sin(phi+angle)+yoffset
    draw.line((int(prevptx), int(prevpty), int(ptx), int(pty)), fill=int(fill), width=int(width))
    width += widthstep
    fill += fillstep if fill < 255 else 0
    prevptx = ptx
    prevpty = pty
    if ptx < 0 or pty < 0 or ptx > im.size[0] or pty > im.size[1]:
      break
    phi+=step

def planet_gen(stars):
  # planet system probability = 0.5
  # number of planets 0-15
  planets = []
  for sid, st in enumerate(stars):
    has_system = random.randrange(0, 6)
    if has_system:
      number_of_planets = random.randrange(1, 15)
      distance = 0
      i = 0
      while i < number_of_planets:
        # gas giant or earth-like ?
        gg = random.randrange(0, 2)
        if gg:
          r = random.randrange(20000, 600000)
          mval = random.randrange(2, 100)
          mexp = random.randrange(26, 28)
          # orbital ellipse params
          distance = distance + random.randrange(1, 200)/100.0
          # orbital angle
          oa = random.gauss(0, 0.7)*120+60
          # orbital velocity (km/s)
          velsign = random.randrange(1, 3, 2)-2
          vo = velsign*random.randrange(5, 50)
          # axial angle
          a = random.gauss(0, 0.7)*120+60
          # equatorial velocity (km/s)
          velsign = random.randrange(1, 3, 2)-2
          ve = velsign*random.randrange(1, 50000)
          
        else:
          r = random.randrange(4000, 10000)
          mval = random.randrange(2, 10)
          mexp = random.randrange(24, 25)
          distance = distance + random.randrange(1, 200)/100.0
          # orbital params
          # orbital angle
          oa = random.gauss(0, 0.7)*120+60
          # orbital velocity (km/s)
          velsign = random.randrange(1, 3, 2)-2
          vo = velsign*random.randrange(5, 50)
          # axial angle
          a = random.gauss(0, 0.7)*120+60
          # equatorial velocity (km/s)
          velsign = random.randrange(1, 3, 2)-2
          ve = velsign*random.randrange(1, 50000)
        planets.append({"r":r, "mval":mval, "mexp":mexp, "distance": distance, "i": oa, "vo": vo, "a": a, "ve": ve, "sid": sid, "pid": i})
        i += 1
  return planets

def moon_gen(planets):
  # moon system rises at farther end of solar system
  # number of moons 0-20
  # at 1 ae probability of moon system existance is 0.3, at 10 ae probability of moon system existence 0.8
  # at 1 ae number of moons is 1 - 2, at 10 ae - 10-20
  # probability threshold = k * distance + b
  # 0.3 = k * 1 + b; b = 0.3 - k
  # 0.6 = k * 10 + b; 0.6 = k * 10 + 0.3 - k; k = 0.3/9
  # max moon count
  # 1 = k * 1 + b
  # 5 = k * 10 + b
  # 1-k=b
  # 5=k*10+1-k
  # 4/9=k
  k_hs=0.3/9.0
  b_hs=0.3-k_hs
  k_mc = 4.0/9.0
  moons = []
  for pid, p in enumerate(planets):
    distance = p["distance"]
    rnd = random.randrange(0, 10)
    has_system = True if rnd>k_hs*distance+b_hs else False
    if has_system:
      number_of_moons = random.randrange(0, int(k_mc*distance)+1)
      moon_distance = 0.0
      i = 0
      while i < number_of_moons:
        # icy or sandy
        icy = random.randrange(0, 2)
        if icy:
          r = random.randrange(200, 1000)
          mval = random.randrange(2, 100)
          mexp = random.randrange(20, 22)
          # orbital ellipse params
          distance = distance + random.randrange(1, 200)/1000.0
          # orbital angle
          oa = random.gauss(0, 0.7)*120+60
          # orbital velocity (km/s)
          velsign = random.randrange(1, 3, 2)-2
          vo = velsign*random.randrange(5, 50)
          # axial angle
          a = random.gauss(0, 0.7)*120+60
          # equatorial velocity (km/s)
          velsign = random.randrange(1, 3, 2)-2
          ve = velsign*random.randrange(1, 50000)
          
        else:
          r = random.randrange(400, 2000)
          mval = random.randrange(2, 10)
          mexp = random.randrange(20, 21)
          distance = distance + random.randrange(1, 200)/1000.0
          # orbital params
          # orbital angle
          oa = random.gauss(0, 0.7)*120+60
          # orbital velocity (km/s)
          velsign = random.randrange(1, 3, 2)-2
          vo = velsign*random.randrange(5, 50)
          # axial angle
          a = random.gauss(0, 0.7)*120+60
          # equatorial velocity (km/s)
          velsign = random.randrange(1, 3, 2)-2
          ve = velsign*random.randrange(1, 50000)
        moons.append({"r":r, "mval":mval, "mexp":mexp, "distance": distance, "i": oa, "vo": vo, "a": a, "ve": ve, "pid": pid})
        i += 1
  return moons

def data_gen():
  im = Image.new("L", (160, 160), 254)
  draw = ImageDraw.Draw(im)
  # main branch
  branch_gen([im, draw], im.size[0]/2.0+random.randrange(0,4), im.size[0]/2.0+random.randrange(0,4), im.size[0]/random.randrange(2,6), -0.1, 0, 5, random.randrange(1, 5), 1.0/random.randrange(10, 100))
  branch_gen([im, draw], im.size[0]/2.0+random.randrange(-1,1), im.size[0]/2.0+random.randrange(-1,1), im.size[0]/random.randrange(2,6),0, 0, 5, random.randrange(1, 5), 1.0/random.randrange(10, 100))
  branch_gen([im, draw], im.size[0]/2.0-random.randrange(0,3), im.size[0]/2.0-random.randrange(0,3), im.size[0]/random.randrange(2,6),0.1, 0, 5, random.randrange(1, 5), 1.0/random.randrange(10, 100))

  # secondary branches
  branch_gen([im, draw], im.size[0]/2.0, im.size[0]/2.0, im.size[0]/random.randrange(2,8), math.pi/3.0+random.randrange(-30,30)*math.pi/180, 0, random.randrange(15,30), random.randrange(1, 3), 1.0/random.randrange(10, 1000))
  branch_gen([im, draw], im.size[0]/2.0, im.size[0]/2.0, im.size[0]/random.randrange(2,8), 2*math.pi/3.0+random.randrange(-30,30)*math.pi/180, 0, random.randrange(15,30), random.randrange(1, 3), 1.0/random.randrange(10, 1000))

  # bar
  ax1 = random.randrange(3, 8)
  ax2 = random.randrange(2, ax1)
  
  draw.ellipse((im.size[0]/2.0-im.size[0]*ax1/48, im.size[0]/2.0-im.size[0]*ax2/48, im.size[0]/2.0+im.size[0]*ax1/48, im.size[0]/2.0+im.size[0]*ax2/48), fill=100)
  del draw
  # im1 = Image.open("top.png")
  # print im1.mode
  im = im.filter(ImageFilter.SMOOTH_MORE)
  im = im.filter(ImageFilter.BLUR)

#  im.save("top1.png", "PNG")

  # im = Image.open("top.png")


  stars = []
  xsize = 2000.0
  ysize = 2000.0
  zsize = 1000.0

  iw, ih = im.size
  sigma = 150
  # x = -1000.0
  # while x<1000:
  #     print ((1/math.sqrt(2.0*math.pi*sigma**2)*math.exp(-(x**2)/(2.0*sigma**2))*100))*zsize/2.0
  #     x+=100
  # exit()


  xscale = xsize/iw
  yscale = ysize/ih
  maxstars = 20000
  nextpoint = maxstars/100.0
  galaxycoloroffset = random.randrange(0.0, 15000.0)
  print "galaxy color offset:", galaxycoloroffset
  while len(stars)<maxstars:
      x = random.gauss(0, 0.7)*xsize/2.0#random.randrange(-xsize/2.0, xsize/2.0)

      y = random.gauss(0, 0.7)*ysize/2.0#random.randrange(-xsize/2.0, xsize/2.0)
      if y**2+x**2>(ysize/2)**2:
          continue

      if random.random() > im.getpixel((int((x+xsize/2.0)/xscale), int((y+ysize/2.0)/yscale)))/255.0:
          z = (random.gauss(0, 0.01))*zsize
          z+=random.random()*z/math.fabs(z)*((1/math.sqrt(2.0*math.pi*sigma**2)*math.exp(-(x**2+y**2)/(2.0*sigma**2))*100))*zsize/2.0
          g = random.gauss(0, 0.1)*60000.0
          temperature = galaxycoloroffset+g if galaxycoloroffset+g > 0 else random.randrange(1000, 1500)
  #random.randrange(2000.0, 60000.0)
          starclass = "O"
          mass = random.randrange(5, 10)/10.0
          r = random.randrange(5, 20000)/10.0
          if temperature >= 1000.0 and temperature < 4500.0:
              starclass = "M"   
          elif temperature < 6000.0:
              starclass = "K"
          elif temperature < 8000.0:
              starclass = "G"
          elif temperature < 10500.0:
              starclass = "F"
          elif temperature < 15000.0:
              starclass = "A"
          elif temperature < 30000.0:
              starclass = "B"

          # if find(lambda e: e["x"] == x and e["y"]==y and e["z"] == z, stars):
          #     continue
          stars.append({"x": x, "y": y, "z": z, "class": starclass, "temperature": temperature, "r": r, "mass": mass, "sid": len(stars)})
          if len(stars)>nextpoint:
              nextpoint += maxstars/100.0
              #print len(stars)/(maxstars/100.0)

  planets = planet_gen(stars)
  moons = moon_gen(planets)
  print "len of stars:", len(stars)
  print "len of planets:", len(planets)
  print "len of moons:", len(moons)
  data = {"stars":stars, "planets": planets, "moons": moons}
  return data
